#!/usr/bin/env python
#
# Copyright 2011 (C) David I. Lehn <dil@lehn.org>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import bisect
import logging

class Range(object):
    """A simple inclusive range between a start and end value."""

    def __init__(self, start_or_range, end=None):
        """Create a range.

        >>> r1 = Range(10,20)
        >>> r2 = Range([10,20])
        >>> r3 = Range(["10","20"])
        >>> assert 10 == r1.start == r2.start == r3.start
        >>> assert 20 == r1.end == r2.end == r3.end
        >>> r3 = Range(10)
        >>> r4 = Range(["10"])
        >>> r5 = Range("10")
        >>> assert 10 == r3.start == r4.start == r5.start
        >>> assert 10 == r3.end == r4.end == r5.end

        """

        if end is None:
            if isinstance(start_or_range, (list, tuple)):
                if len(start_or_range) == 1:
                    self.start = self.end = int(start_or_range[0])
                else:
                    self.start, self.end = [int(v) for v in start_or_range]
            else:
                self.start = self.end = int(start_or_range)
        else:
            self.start = int(start_or_range)
            self.end = int(end)

    def __cmp__(self, other):
        """Compare ranges for use in sorting ranges.
        
        >>> r10_20 = Range(10,20)
        >>> r20_30 = Range(20,30)
        >>> r30_40 = Range(30,40)
        >>> r10_15 = Range(10,15)
        >>> cmp(r10_20, r10_20)
        0
        >>> cmp(r10_20, r30_40)
        -1
        >>> cmp(r30_40, r10_20)
        1
        >>> cmp(r10_20, r20_30)
        -1
        """

        if self.start == other.start and self.end == other.end:
            return 0
        elif self.start > (other.end + 1):
            return 1
        else:
            return -1


def _str_range(r):
    """Utility function to print a single range.
    
    >>> _str_range([1,1])
    '1'
    >>> _str_range([1,2])
    '1-2'

    """
    if r[0] == r[1]:
        return "%d" % (r[0])
    else:
        return "%d-%d" % (r[0], r[1])


class SparseRange(object):
    """Create a SparseRange.

    SparseRanges have the following features:

    - Keep track of any number of ranges at once.
    - Add and remove arbitrary overlapping ranges.
    - Add and return a sized range not in the current object.
    - Remove and return a sized range in the current object.

    """
    def __init__(self, range=None):
        """Initialize a SparseRange.

        >>> sr = SparseRange()
        >>> print str(sr)
        <BLANKLINE>
        >>> sr = SparseRange((10,20))
        >>> print str(sr)
        10-20
        >>> sr = SparseRange((10,20))
        >>> print str(sr)
        10-20

        """
        self.clear()
        if range is not None:
            self.add_range(range)

    def __str__(self):
        return ",".join([_str_range(r) for r in self.ranges()])

    def __len__(self):
        """Get the total number of values in this sparse range.

        >>> sr = SparseRange()
        >>> sr.add_range([1,2])
        >>> len(sr)
        2
        >>> sr.add_range([10,20])
        >>> len(sr)
        13

        """
        return sum([end - start + 1 for start, end in self.ranges()])

    def clear(self):
        """Remove all ranges.

        >>> sr = SparseRange()
        >>> sr.add_range([1,2])
        >>> sr.clear()
        >>> sr.ranges()
        []

        """
        self.starts = []
        self.stops = []

    def ranges(self):
        """Return a sorted list of (start, end) pairs.
        
        >>> sr = SparseRange()
        >>> sr.add_range([1,2])
        >>> sr.add_range([5,6])
        >>> sr.ranges()
        [(1, 2), (5, 6)]

        """

        return zip(self.starts, self.stops)

    def add_range(self, range):
        """Add a range or ranges.
        
        The range can be a Range instance, a single index, a (start, end) pair,
        a single (index), or a string as returned by str on a SparseRange.

        The new range is allowed to overlap current ranges.

        >>> sr = SparseRange()
        >>> sr.add_range([10,20])
        >>> sr.add_range([30,40])
        >>> print str(sr)
        10-20,30-40
        >>> sr.add_range(25)
        >>> print str(sr)
        10-20,25,30-40

        Adding a range with duplicate values:
        >>> sr = SparseRange()
        >>> sr.add_range([10,20])
        >>> sr.add_range([14,16])
        >>> print str(sr)
        10-20

        Adding a range will combine ranges:
        >>> sr = SparseRange()
        >>> sr.add_range([10,20])
        >>> sr.add_range([21,30])
        >>> print str(sr)
        10-30

        Adding a range with duplicate and new values:
        >>> sr = SparseRange()
        >>> sr.add_range([10,20])
        >>> sr.add_range([15,25])
        >>> print str(sr)
        10-25
        >>> sr.add_range([5,35])
        >>> print str(sr)
        5-35
        >>> sr.add_range([40,45])
        >>> sr.add_range([50,55])
        >>> print str(sr)
        5-35,40-45,50-55
        >>> sr.add_range([5,65])
        >>> print str(sr)
        5-65

        Flexible range formats:
        >>> sr = SparseRange()
        >>> sr.add_range(10)
        >>> sr.add_range(Range(20,25))
        >>> sr.add_range("30")
        >>> sr.add_range([40,45])
        >>> sr.add_range("50-55")
        >>> print str(sr)
        10,20-25,30,40-45,50-55

        """
        if isinstance(range, Range):
            # already a range
            pass
        elif isinstance(range, (int, list, tuple)):
            # convert simple ranges
            range = Range(range)
        elif isinstance(range, SparseRange):
            # add sub-ranges
            for r in range.ranges():
                self.add_range(r)
            return
        elif isinstance(range, basestring):
            # parse and add sub-ranges
            for r in range.split(","):
                self.add_range(r.split("-"))
            return

        rl = len(self.starts)
        if rl == 0:
            # special case when empty
            self.starts.append(range.start)
            self.stops.append(range.end)
        else:
            spos = bisect.bisect_left(self.starts, range.start)
            # check if next range has same start
            if (spos + 1) < rl and self.starts[spos + 1] == range.start:
                spos += 1
            # check if previous range can be combined
            elif (spos - 1) >= 0 and self.stops[spos - 1] >= (range.start - 1):
                spos -= 1
                range.start = self.starts[spos]

            epos = bisect.bisect_right(self.stops, range.end)
            # check if previous range has same end
            if (epos - 1) >= 0 and self.stops[epos - 1] == range.end:
                epos -= 1
            # check if next range can be combined
            elif epos < rl and self.starts[epos] <= (range.end + 1):
                range.end = self.stops[epos]
                epos += 1

            self.starts[spos:epos] = [range.start]
            self.stops[spos:epos] = [range.end]

    def add_count(self, count):
        """Add and return count indexes to the end of a non-empty range.
        
        Params:
            count: number of indexes to add

        Return:
            a SparseRange representing the added range

        >>> sr = SparseRange()
        >>> sr.add_range(0)
        >>> sr.add_count(1).ranges()
        [(1, 1)]
        >>> sr.ranges()
        [(0, 1)]
        >>> sr = SparseRange()
        >>> sr.add_range([0,10])
        >>> sr.add_count(10).ranges()
        [(11, 20)]
        >>> sr.ranges()
        [(0, 20)]

        """
        start = self.stops[-1] + 1
        self.stops[-1] += count
        return SparseRange([start, start + count - 1])

    def remove_range(self, range, partial=False):
        """Remove and return the specified range.

        Params:
            partial: if True may return a partial range.
            partial: if True may return a partial range, if False will raise
                IndexError if range can not be removed.

        Return:
            A SparseRange with the removed indexes.
        """

    def remove_count(self, count, partial=False):
        """Remove and return count indexes.

        Params:
            partial: if True may return less than count indexes, if False will
                raise IndexError if range if size count can not be removed.

        Return:
            A SparseRange with the removed indexes.

        >>> sr = SparseRange()
        >>> sr.add_range([1,10])
        >>> r = sr.remove_count(1)
        >>> r.ranges()
        [(1, 1)]
        >>> sr.ranges()
        [(2, 10)]

        >>> r = sr.remove_count(9)
        >>> r.ranges()
        [(2, 10)]
        >>> sr.ranges()
        []

        >>> r = sr.remove_count(1)
        Traceback (most recent call last):
           ...
        IndexError: Range too small

        """
        removed = SparseRange()
        # avoid repeated length calls
        srlen = len(self)
        if srlen < count and partial == False:
            raise IndexError("Range too small")
        rlen = 0
        while srlen > 0 and rlen < count:
            # next range length
            next_len = self.stops[0] - self.starts[0] + 1
            # amount to remove from this range
            next_rem = min(next_len, count - rlen)
            # adjust lengths
            rlen += next_rem
            srlen -= next_rem
            if next_len == next_rem:
                # return and remove this whole range
                removed.add_range([self.starts[0], self.stops[0]])
                del self.starts[0]
                del self.stops[0]
            else:
                # return and remove a partial range
                start = self.starts[0]
                stop = start + next_rem - 1
                removed.add_range([start, stop])
                self.starts[0] = start + next_rem
        return removed

if __name__ == "__main__":
    import doctest
    doctest.testmod()
